#pragma once

#include <map>
#include "../VulkanRHI.h"
#include "../../Utils/FileUtil.h"
#include "../../Templates/RefCounting.h"
#include "../RHIResource.h"
#include "../../Templates/NewObject.h"
#include "ShaderMacros.h"
#include "../Texture/RHISampler.h"
#include "../Texture/RHITexture.h"
#include "../Buffer/BufferView.h"
#include "../Buffer/RHIBuffer.h"
#include "../Raytracing/AccelerationStructure.h"
#include "../Texture/RenderTarget.h"


class RHIDescriptor;
class ShaderConstant;
class ShaderParam;
class ShaderAdapter;

class RHIShader
{
	struct VertexInputBinding
	{
		VkVertexInputBindingDescription mBinding;
		std::unordered_map<std::string, VkVertexInputAttributeDescription> mAttributes;
		VertexInputBinding()
		{
			mBinding.binding = UINT32_MAX;
			mBinding.stride = UINT32_MAX;
			mBinding.inputRate = VK_VERTEX_INPUT_RATE_MAX_ENUM;
		}

		VertexInputBinding(std::vector<std::string>& attributes, uint32 binging, uint32 stride, VkVertexInputRate inputRate)
		{
			mBinding.binding = binging;
			mBinding.stride = stride;
			mBinding.inputRate = inputRate;

			for (const auto& attribute : attributes)
			{
				VkVertexInputAttributeDescription attributeDescription;
				attributeDescription.location = UINT32_MAX;
				attributeDescription.binding = binging;
				attributeDescription.format = VK_FORMAT_UNDEFINED;
				attributeDescription.offset = UINT32_MAX;
				mAttributes.insert({ attribute, attributeDescription });
			}
		}

		VertexInputBinding(std::vector<std::pair<std::string, uint32>>& attributes, uint32 binging, uint32 stride, VkVertexInputRate inputRate)
		{
			mBinding.binding = binging;
			mBinding.stride = stride;
			mBinding.inputRate = inputRate;

			for (const auto& attribute : attributes)
			{
				VkVertexInputAttributeDescription attributeDescription;
				attributeDescription.location = UINT32_MAX;
				attributeDescription.binding = binging;
				attributeDescription.format = VK_FORMAT_UNDEFINED;
				attributeDescription.offset = attribute.second;
				mAttributes.insert({ attribute.first, attributeDescription });
			}
		}

		VertexInputBinding(std::vector<std::pair<std::string, std::pair<uint32, VkFormat>>>& attributes, uint32 binging, uint32 stride, VkVertexInputRate inputRate)
		{
			mBinding.binding = binging;
			mBinding.stride = stride;
			mBinding.inputRate = inputRate;

			for (const auto& attribute : attributes)
			{
				VkVertexInputAttributeDescription attributeDescription;
				attributeDescription.location = UINT32_MAX;
				attributeDescription.binding = binging;
				attributeDescription.format = attribute.second.second;
				attributeDescription.offset = attribute.second.first;
				mAttributes.insert({ attribute.first, attributeDescription });
			}
		}
	};


public:

	RHIShader(std::string path);

	~RHIShader();

	VkPipelineShaderStageCreateInfo GetShaderStageCreateInfo() const;

	VkPipelineVertexInputStateCreateInfo& GetVertexInputStateCreateInfo();

	VkPipelineInputAssemblyStateCreateInfo& GetInputAssemblyStateCreateInfo();

	VkShaderStageFlagBits GetShaderStage() const;

	VkPushConstantRange GetPushConstantRange() const;

	uint32 GetHash() const;

protected:

	void CompileShader();

	virtual void CollectShaderInfos() {};

	// Vertex input state

	void GROUP_VERTEX_BINGING(std::vector<std::string> attributes, VkVertexInputRate inputRate = VK_VERTEX_INPUT_RATE_MAX_ENUM, uint32 stride = UINT32_MAX);

	void GROUP_VERTEX_BINGING_OFFSET(std::vector<std::pair<std::string, uint32>> attributes, VkVertexInputRate inputRate = VK_VERTEX_INPUT_RATE_MAX_ENUM, uint32 stride = UINT32_MAX);

	void GROUP_VERTEX_BINGING_OFFSET_FORMAT(std::vector<std::pair<std::string, std::pair<uint32, VkFormat>>> attributes, VkVertexInputRate inputRate = VK_VERTEX_INPUT_RATE_MAX_ENUM, uint32 stride = UINT32_MAX);

	// Input assembly state
	void SET_INPUT_ASSEMBLY(VkPrimitiveTopology topology, VkBool32 primitiveRestartEnable = VK_FALSE);

	// Push constant

	void PUSH_CONSTANT_OFFSET(uint32 offset);

	// Constant id

	uint32 AddSpecializationMapEntry(VkDeviceSize size);

	uint8* GetSpecializationData(uint32 constantId);

	// Shader param
	void AddShaderParam(const std::string& name, ShaderParam* param);

	// Set adapter

	void SetShaderAdapter(ShaderAdapter* shaderAdapter);

	RHIDescriptor* GetDescriptor() const;

protected:

	std::string mPath;
	VkShaderStageFlagBits mShaderStage;
	ShaderAdapter* mShaderAdapter;
	VkShaderModule mShaderModule;
	VkPipelineShaderStageCreateInfo mShaderStageCreateInfo;


	std::vector<VertexInputBinding> mVertexInputTables;
	std::vector<VkVertexInputBindingDescription> mVertexInputBindingDescriptions;
	std::vector<VkVertexInputAttributeDescription> mVertexInputAttributeDescriptions;
	VkPipelineVertexInputStateCreateInfo mVertexInputStateCreateInfo;

	VkPipelineInputAssemblyStateCreateInfo mInputAssemblyStateCreateInfo;

	VkPushConstantRange mPushConstantRange;

	VkSpecializationInfo mSpecializationInfo;
	std::vector<VkSpecializationMapEntry> mSpecializationMapEntrys;

	std::unordered_map<std::string, ShaderParam*> mShaderParams;

	uint32 mHash;

	friend class ShaderMap;
	friend class ShaderAdapter;
	friend class ShaderConstant;
	friend class ShaderParam;
	friend class RHIDescriptor;

	template<typename ShaderType>
	friend class ShaderPtr;
};


class ShaderConstant
{
	RHIShader* mShader;
	uint32 mConstantId;
public:
	ShaderConstant(RHIShader* shader, uint32 size) : mShader(shader)
	{
		mConstantId = mShader->AddSpecializationMapEntry(size);
	}
protected:
	uint8* GetData()
	{
		return mShader->GetSpecializationData(mConstantId);
	}
};

class ShaderParam
{
	friend class RHIShader;
	friend class RHIDescriptor;

protected:

	RHIShader* mShader;
	VkDescriptorType mType;
	uint32 mSet;
	uint32 mBinging;
	uint32 mCount;

protected:

	void SetSAMPLER(RefCountPtr<RHISampler> sampler);

	void SetSAMPLERs(const std::vector<RefCountPtr<RHISampler>>& samplers);

	void SetCOMBINED_IMAGE_SAMPLER(RefCountPtr<RHITexture> texture, RefCountPtr<RHISampler> sampler, VkImageLayout layout);

	void SetCOMBINED_IMAGE_SAMPLERs(const std::vector<RefCountPtr<RHITexture>>& textures, const std::vector<RefCountPtr<RHISampler>>& samplers, const std::vector<VkImageLayout>& layouts);

	void SetSAMPLED_IMAGE(RefCountPtr<RHITexture> texture, VkImageLayout layout);

	void SetSAMPLED_IMAGEs(const std::vector<RefCountPtr<RHITexture>>& textures, const std::vector<VkImageLayout>& layouts);

	void SetSTORAGE_IMAGE(RefCountPtr<RHITexture> texture, VkImageLayout layout);

	void SetSTORAGE_IMAGEs(const std::vector<RefCountPtr<RHITexture>>& textures, const std::vector<VkImageLayout>& layouts);

	void SetINPUT_ATTACHMENT(RefCountPtr<RenderTarget> texture, VkImageLayout layout);

	void SetUNIFORM_TEXEL_BUFFER(RefCountPtr<BufferView> bufferView);

	void SetUNIFORM_TEXEL_BUFFERs(const std::vector<RefCountPtr<BufferView>>& bufferViews);

	void SetSTORAGE_TEXEL_BUFFER(RefCountPtr<BufferView> bufferView);

	void SetSTORAGE_TEXEL_BUFFERs(const std::vector<RefCountPtr<BufferView>>& bufferViews);

	void SetUNIFORM_BUFFER(RefCountPtr<RHIBuffer> buffer, VkDeviceSize offset, VkDeviceSize range);

	void SetUNIFORM_BUFFERs(const std::vector<RefCountPtr<RHIBuffer>>& buffers, const std::vector<VkDeviceSize>& offsets, const std::vector<VkDeviceSize>& ranges);

	void SetSTORAGE_BUFFER(RefCountPtr<RHIBuffer> buffer, VkDeviceSize offset, VkDeviceSize range);

	void SetSTORAGE_BUFFERs(const std::vector<RefCountPtr<RHIBuffer>>& buffers, const std::vector<VkDeviceSize>& offsets, const std::vector<VkDeviceSize>& ranges);

	void SetUNIFORM_BUFFER_DYNAMIC(RefCountPtr<RHIBuffer> buffer, VkDeviceSize offset, VkDeviceSize range);

	void SetUNIFORM_BUFFER_DYNAMICs(const std::vector<RefCountPtr<RHIBuffer>>& buffers, const std::vector<VkDeviceSize>& offsets, const std::vector<VkDeviceSize>& ranges);

	void SetSTORAGE_BUFFER_DYNAMIC(RefCountPtr<RHIBuffer> buffer, VkDeviceSize offset, VkDeviceSize range);

	void SetSTORAGE_BUFFER_DYNAMICs(const std::vector<RefCountPtr<RHIBuffer>>& buffers, const std::vector<VkDeviceSize>& offsets, const std::vector<VkDeviceSize>& ranges);

	void SetINLINE_UNIFORM_BLOCK_EXT(void* value);

	void SetACCELERATION_STRUCTURE_KHR(RefCountPtr<AccelerationStructure> topLevelAS);

	void SetACCELERATION_STRUCTURE_KHRs(const std::vector<RefCountPtr<AccelerationStructure>>& topLevelASs);
	
public:

	ShaderParam(RHIShader* shader, VkDescriptorType type, const char* name) :
		mShader(shader),
		mType(type),
		mSet(UINT32_MAX),
		mBinging(UINT32_MAX),
		mCount(UINT32_MAX)
	{
		mShader->AddShaderParam(name, this);
	}
};


class ShaderAdapter
{
public:
	ShaderAdapter();
	virtual ~ShaderAdapter();
	uint32 AddRef() const
	{
		return uint32(++NumRefs);
	}
	uint32 Release() const
	{
		uint32 Refs = uint32(--NumRefs);
		return Refs;
	}
	uint32 GetRefCount() const
	{
		return uint32(NumRefs);
	}
	void SetShader(RHIShader* shader);
	void SetDescriptor(RHIDescriptor* descriptor);
private:
	mutable int32 NumRefs;
public:
	RHIShader* mShader;
	void* mSpecializationData;
	RHIDescriptor* mDescriptor;
};


#if 0

void SetPushConstant(uint32 size, uint32 offset);

void RHIShader::SetPushConstant(uint32 size, uint32 offset)
{
	mPushConstantRange.stageFlags = mShaderStage;
	mPushConstantRange.size = size;
	mPushConstantRange.offset = offset;
}

#endif